package org.javaexpert.spaces.masterworker;

import net.jini.core.entry.Entry;
import net.jini.core.entry.UnusableEntryException;
import net.jini.core.lease.Lease;
import static net.jini.core.lease.Lease.FOREVER;
import net.jini.core.lease.LeaseDeniedException;
import net.jini.core.lease.LeaseException;
import net.jini.core.transaction.Transaction;
import net.jini.core.transaction.TransactionException;
import net.jini.space.JavaSpace;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.javaexpert.spaces.util.CallableInTxn;
import org.javaexpert.spaces.util.CriticalSpaceException;
import org.javaexpert.spaces.util.RunnableInTxn;
import org.javaexpert.spaces.util.TransactionTemplate;
import sun.reflect.generics.reflectiveObjects.NotImplementedException;

import static java.lang.String.format;
import java.rmi.RemoteException;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

public class TaskBag {
    String id;
    JavaSpace space;
    TransactionTemplate txnTemplate;
    TaskMonitor taskMonitor;
    long TASK_LEASE_TIME = WEEK;

    Entry taskBagTemplate;
    private static final Long MINUTE = 1000 * 60L;
    private static final Log log = LogFactory.getLog("main");

    public TaskBag(String id, JavaSpace space, TransactionTemplate transactionTemplate) throws RemoteException {
        assert space != null;
        assert transactionTemplate != null;
        assert id != null;
        this.id = id;
        this.space = space;
        this.txnTemplate = transactionTemplate;
        taskMonitor = new TaskMonitor(space, txnTemplate, id);
        taskBagTemplate = new TaskBagEntry(id);
    }

    static final int WEEK = 1000 * 3600 * 24 * 7;

    public <T> Future<T> submit(final Callable<T> callable) throws RemoteException, LeaseDeniedException, TransactionException, UnusableEntryException, InterruptedException {
        return submit(callable, 3);
    }

    public TaskGroup submit(List<Callable> tasks) throws TransactionException, LeaseDeniedException, UnusableEntryException, InterruptedException, RemoteException {
        Long groupId = generateNewGroupId();
        for (Callable c : tasks) {
            submit(c, 3, groupId);
        }
        return new TaskBagTaskGroup(groupId);
    }

    public <T> Future<T> submit(final Callable<T> callable, final int maxRetryCount) throws RemoteException, LeaseDeniedException, TransactionException, UnusableEntryException, InterruptedException {
        assert callable != null;
        return submit(callable, maxRetryCount, null);
    }

    protected <T> Future<T> submit(final Callable<T> callable, final int maxRetryCount, final Long groupId) throws RemoteException, LeaseDeniedException, TransactionException, UnusableEntryException, InterruptedException {
        assert callable != null;

        final Long taskId = generateNewTaskId();
        System.out.println("taskid=" + taskId);
        txnTemplate.runInTxn(new RunnableInTxn() {
            public void run(Transaction txn) throws Exception {
                space.write(new TaskEntry(taskId, callable, maxRetryCount), txn, TASK_LEASE_TIME);
                taskMonitor.addNewTaskStatus(taskId, groupId, txn);
            }
        });
        return new TaskBagFuture<T>(taskId, callable, space, this);
    }

    public void acceptAndRunTasks(final String workerId) throws CriticalSpaceException {
        for (; ;) {
            try {
                space.read(new TaskEntry(), null, 10 * 1000);
                txnTemplate.runInTxn(10 * 1000L, 60 * 24 * MINUTE, new TryToGetAndExecuteTask(workerId));
            } catch (RuntimeException re) {
                ignoreNonCriticalExceptions(re.getCause());
            } catch (Exception e) {
                ignoreNonCriticalExceptions(e);
            }
        }

    }

    public <T> T consumeTaskResult(Long taskId) throws ExecutionException {
        TaskBag.TaskResultEntry template = new TaskBag.TaskResultEntry(taskId);
        TaskBag.TaskResultEntry<T> resultEntry;
        Transaction txn = txnTemplate.getTransaction(Lease.FOREVER);
        try {
            resultEntry = (TaskBag.TaskResultEntry<T>) space.take(template, txn, Lease.FOREVER);
            taskMonitor.consume(resultEntry.id, txn);
            txn.commit();
        } catch (Exception e) {
            txnTemplate.abort(txn);
            throw new ExecutionException(format("Error while getting result for taskId=%s", taskId), e);
        }
        if (resultEntry instanceof TaskBag.TaskFailedResultEntry) {
            throw new ExecutionException(((TaskBag.TaskFailedResultEntry) resultEntry).errorMessage, null);
        }
        return resultEntry.result;
    }

    private void ignoreNonCriticalExceptions(Throwable e) throws CriticalSpaceException {
        if (e instanceof LeaseException || e instanceof TransactionException) {
            ignoreException(e);
        }
//        if(        e instanceof InterruptedException
//                || e instanceof RemoteException
//                || e instanceof UnusableEntryException
//                ) {
        throw new CriticalSpaceException(e);
//        }
    }


    private Long generateNewTaskId() {
        return txnTemplate.runInTxn(new CallableInTxn<Long>() {
            public Long call(Transaction txn) throws Exception {
                TaskBagEntry taskBag = takeOrCreateTaskBagEntry(txn);
                Long taskId = taskBag.incrementTaskCounter();
                space.write(taskBag, txn, FOREVER);
                return taskId;
            }
        });
    }

    private Long generateNewGroupId() {
        return txnTemplate.runInTxn(new CallableInTxn<Long>() {
            public Long call(Transaction txn) throws Exception {
                TaskBagEntry taskBag = takeOrCreateTaskBagEntry(txn);
                Long groupId = taskBag.incrementGroupCounter();
                space.write(taskBag, txn, FOREVER);
                return groupId;
            }
        });
    }

    private TaskBagEntry takeOrCreateTaskBagEntry(Transaction txn) throws UnusableEntryException, TransactionException, InterruptedException, RemoteException {
        TaskBagEntry taskBag = (TaskBagEntry) space.takeIfExists(taskBagTemplate, txn, FOREVER);
        if (taskBag == null) {
            taskBag = new TaskBagEntry(id).initCounters();
        }
        return taskBag;
    }


    public Future submit(Runnable runnable) {
        throw new NotImplementedException();
    }

    public static class TaskEntry<T> implements Entry {
        public Long id;
        public Callable<T> callable;
        public Integer maxRetryCount;

        public TaskEntry() {
        }

        public TaskEntry(Long id, Callable<T> callable, Integer maxRetryCount) {
            this.id = id;
            this.callable = callable;
            this.maxRetryCount = maxRetryCount;
        }


        public String toString() {
            return "TaskEntry{" +
                    "id=" + id +
                    ", callable=" + callable +
                    ", maxRetryCount=" + maxRetryCount +
                    '}';
        }

        public TaskResultEntry<T> createResultEntry(T res) {
            return new TaskResultEntry<T>(id, res);
        }

        private TaskFailedResultEntry createFailedResultEntry(final String comment) {
            return new TaskFailedResultEntry(id, comment);
        }

        private FailedTaskEntry createFailedTaskEntry() {
            return new FailedTaskEntry(id, this);
        }
    }

    public static class TaskResultEntry<T> implements Entry {
        public Long id;
        public T result;


        public TaskResultEntry() {
        }

        public TaskResultEntry(Long id) {
            this.id = id;
        }

        public TaskResultEntry(Long id, T result) {
            this.id = id;
            this.result = result;
        }
    }

    public static class TaskFailedResultEntry extends TaskResultEntry {
        public String errorMessage;

        public TaskFailedResultEntry() {
        }

        public TaskFailedResultEntry(Long id, String errorMessage) {
            super(id, null);
            this.errorMessage = errorMessage;
        }
    }

    public static class TaskBagEntry implements Entry {
        public String id;
        public Long taskCounter;
        public Long groupCounter;

        public TaskBagEntry() {
        }

        public TaskBagEntry(String id) {
            this.id = id;
        }

        protected TaskBagEntry initCounters() {
            this.taskCounter = (id.hashCode() % 10000) * 1000L * 1000L * 1000L;
            this.groupCounter = 0L;
            return this;
        }

        public long incrementTaskCounter() {
            return ++taskCounter;
        }

        public long incrementGroupCounter() {
            return ++groupCounter;
        }
    }

    public static class TaskLogEntry implements Entry {
        public Long taskId;
        public Long index;

        public String txt;

    }

    public static class FailedTaskEntry implements Entry {
        public TaskEntry task;
        public Long taskId;

        public FailedTaskEntry() {
        }

        public FailedTaskEntry(Long id, TaskEntry task) {
            this.task = task;
            this.taskId = id;
        }
    }

    public static void ignoreException(Throwable e) {
        System.err.println("Ignoring the following exception");
        e.printStackTrace();
    }

    private class TryToGetAndExecuteTask implements RunnableInTxn {
        private final String workerId;

        public TryToGetAndExecuteTask(String workerId) {
            this.workerId = workerId;
        }

        public void run(Transaction txn) throws Exception, RemoteException, UnusableEntryException, TransactionException, InterruptedException {
            TaskEntry task = (TaskEntry) space.takeIfExists(new TaskEntry(), txn, 0);

            if (!validateTask(txn, task)) {
                return;
            }

            taskMonitor.startWork(task.id, workerId);
            try {
                Object res;
                try {
                    log.info("Running " + task);
                    res = task.callable.call();
                    log.info("Task OK");
                } catch (Exception e) {
                    log.info("Task failed", e);
                    throw new TaskExecutionException(e);
                }

                space.write(task.createResultEntry(res), txn, FOREVER);
                taskMonitor.succeed(task.id, txn);
            } catch (Exception e) {
                try {
                    TaskStatusData taskStatus = taskMonitor.failOrRestart(task.id, e.getMessage(), task.maxRetryCount);
                    if (taskStatus.isFailed()) {
                        markTaskEntryAsFailed(task, txn, taskStatus.comment);
                    }
                } catch (Exception e1) {
                    log.error("Marking the task as 'failed' failed.", e1);
                }

                throw e;
            }
        }

        private boolean validateTask(Transaction txn, TaskEntry task) throws RemoteException, UnusableEntryException, TransactionException, InterruptedException {
            if (task == null) {
                return false;
            }

            TaskStatusData taskStatus = taskMonitor.getTaskStatus(task.id);
            if (taskStatus.isBeingExecuted()) {
                log.warn("Found task with status " + taskStatus.status + ". Marking it as failed and retrying, if eligable for retry.\n " + taskStatus.toString());
                taskStatus = taskMonitor.failOrRestart(task.id, "Probably worker " + taskStatus.workerId + " died", task.maxRetryCount);
            }

            if (taskStatus.isFailed()) {
                markTaskEntryAsFailed(task, txn, taskStatus.comment);
                return false;
            }
            return true;
        }

        private void markTaskEntryAsFailed(TaskEntry task, Transaction txn, final String comment) throws TransactionException, RemoteException {
            space.write(task.createFailedTaskEntry(), txn, TASK_LEASE_TIME);
            space.write(task.createFailedResultEntry(comment), txn, TASK_LEASE_TIME);
        }

    }

    private class TaskBagTaskGroup implements TaskGroup {
        private Long groupId;

        public TaskBagTaskGroup(Long groupId) {
            this.groupId = groupId;
        }

        public Object take() throws ExecutionException {
            try {
                Long taskId = taskMonitor.waitForFinishedTaskForGroup(groupId, Lease.FOREVER);
                return consumeTaskResult(taskId);
            } catch (Exception e) {
                throw new ExecutionException(e);
            }
        }

        public void cancel() {
            throw new NotImplementedException();
        }

        public Long getId() {
            return groupId;
        }
    }
}
